{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2afd7255",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12e79ccb",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7dfd417d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exporti\n",
    "import io\n",
    "import time\n",
    "import random\n",
    "from pathlib import Path\n",
    "\n",
    "from fastprogress import progress_bar, master_bar\n",
    "import fastprogress\n",
    "\n",
    "import numpy as np\n",
    "import pylab as plt\n",
    "import math\n",
    "\n",
    "import IPython\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data.dataloader import DataLoader\n",
    "from torch.profiler import record_function\n",
    "\n",
    "import webdataset as wds\n",
    "\n",
    "torch.backends.cudnn.benchmark = True\n",
    "torch.backends.cudnn.enabled = True\n",
    "torch.backends.cuda.matmul.allow_tf32 = True\n",
    "torch.set_float32_matmul_precision('medium')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "692f2f8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class SimpleVisual:\n",
    "    def __init__ (self, model, masterbar, total_steps):\n",
    "        self.model = model\n",
    "        self.masterbar = masterbar\n",
    "        self.total_steps = total_steps\n",
    "        self.epochs = total_steps // masterbar.main_bar.total\n",
    "        \n",
    "        gs = plt.GridSpec(2, 1, height_ratios=[3,1])\n",
    "        graph_fig = plt.figure(figsize=(10,6))\n",
    "        self.graph_fig = graph_fig\n",
    "        self.loss_p = graph_fig.add_subplot(gs[0])\n",
    "        self.lr_p = graph_fig.add_subplot(gs[1], sharex=self.loss_p)\n",
    "        self.lr_p.tick_params('x', labelbottom=False)\n",
    "        self.graph_out = None\n",
    "        \n",
    "        self.its = []\n",
    "        self.train_losses = []\n",
    "        self.val_losses = []\n",
    "        self.lr_history = []\n",
    "            \n",
    "    def show(self):\n",
    "        self.start_t = time.time()\n",
    "        self.masterbar.write([\"samples\", \"train\", \"val\", \"time\"], table=True)\n",
    "        self.graph_out = display(self.graph_fig, display_id=True, clear=True)\n",
    "    \n",
    "    def hide(self):\n",
    "        if self.graph_out is not None:\n",
    "            self.graph_out.update(IPython.display.HTML(''))\n",
    "    \n",
    "    def plot(self):\n",
    "        loss_p, lr_p = self.loss_p, self.lr_p\n",
    "        loss_p.clear()\n",
    "        loss_p.plot(self.its, self.train_losses)\n",
    "        loss_p.plot(self.its, self.val_losses)\n",
    "        loss_p.set_xlim(0, self.total_steps)\n",
    "        loss_p.set_yscale('log')\n",
    "        lr_p.clear()\n",
    "        lrs = np.array(self.lr_history)\n",
    "        lr_p.plot(self.its, lrs)\n",
    "        self.graph_out.update(self.graph_fig)\n",
    "    \n",
    "    def add_data(self, it, lr, train_loss, val_los):\n",
    "        self.its.append(it)\n",
    "        self.train_losses.append(train_loss)\n",
    "        self.val_losses.append(val_los)\n",
    "        self.lr_history.append(lr)\n",
    "        self.plot()\n",
    "\n",
    "    def add_table_row(self, it, avg_train_loss, val_loss):\n",
    "        elapsed_t = time.time() - self.start_t\n",
    "        self.masterbar.write([it, f\"{avg_train_loss:.5f}\", f\"{val_loss:.5f}\", fastprogress.core.format_time(elapsed_t)], table=True)\n",
    "    \n",
    "    def on_iter(self, bar, it, avg_train_loss, val_loss):\n",
    "        epoch = math.ceil(it / self.total_steps * self.epochs)\n",
    "        bar.comment = f\"#{epoch}/{self.epochs} loss: {avg_train_loss:.3f} / {val_loss:.3f}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac94bf44",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "\n",
    "# FIXME: we need to keep this synchronised with the validation code below...\n",
    "def validate(model, val, half=True, bs=16, drop_last=False, dl_workers=8, device=\"cuda\"):\n",
    "    if isinstance(val, torch.utils.data.IterableDataset):\n",
    "        val_loader = wds.WebLoader(val, batch_size=None, num_workers=dl_workers, drop_last=drop_last) \\\n",
    "            .unbatched().shuffle(1024).batched(bs)\n",
    "    else:\n",
    "        val_loader = DataLoader(val, batch_size=bs, num_workers=dl_workers, pin_memory=True, drop_last=drop_last)\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        val_loss = 0\n",
    "        val_samples = 0\n",
    "        for args in val_loader:\n",
    "            args = [x.to(device, non_blocking=True) for x in args]\n",
    "            with torch.autocast(device_type=device, dtype=torch.float16 if half else torch.float32, enabled=device!='cpu'):\n",
    "                ps, loss = model(*args)\n",
    "            N = args[0].shape[0]\n",
    "            val_loss += loss.mean().item() * N\n",
    "            val_samples += N\n",
    "        val_loss = val_loss / val_samples\n",
    "\n",
    "    return val_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9da7d029",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "def train(checkpoint_path, model, train, val, half=True, bs=16, lr=1e-4, drop_last=False,\n",
    "          weight_decay=0.1, warmup_steps=10000, epochs=10, clip_gradient_norm=None,\n",
    "          dl_workers=8, visual_class = SimpleVisual, profiler=None,\n",
    "          run_valid_every_iters=8000, table_row_every_iters=80000, chkpt_every_iters=None,\n",
    "          device=\"cuda\", trainable_params=None):\n",
    "    if chkpt_every_iters is None:\n",
    "        chkpt_every_iters = table_row_every_iters\n",
    "    \n",
    "    mb = master_bar(range(epochs))\n",
    "    if isinstance(train, torch.utils.data.IterableDataset):\n",
    "        pct_start = min(0.3, warmup_steps / (epochs * (train.total_samples//bs)))\n",
    "        visual = visual_class(model, mb, epochs * train.total_samples)\n",
    "#         pct_start = min(0.3, warmup_steps / (epochs * len(train)))\n",
    "#         visual = visual_class(model, mb, epochs*len(train)*bs)\n",
    "    else:\n",
    "        pct_start = min(0.3, warmup_steps / (epochs * len(train) / bs))\n",
    "        visual = visual_class(model, mb, epochs*len(train))\n",
    "    model.visual = visual\n",
    "    \n",
    "    Path(checkpoint_path).mkdir(exist_ok=True)\n",
    "\n",
    "    if isinstance(train, torch.utils.data.IterableDataset):\n",
    "#         train_loader = DataLoader(train, batch_size=None, num_workers=dl_workers, pin_memory=True, drop_last=False, shuffle=False)\n",
    "#         val_loader = DataLoader(val, batch_size=None, num_workers=dl_workers, pin_memory=True, drop_last=False)\n",
    "        train_loader = wds.WebLoader(train, batch_size=None, num_workers=dl_workers, drop_last=drop_last) \\\n",
    "            .unbatched().shuffle(1024).batched(bs, partial=False)\n",
    "        val_loader = wds.WebLoader(val, batch_size=None, num_workers=dl_workers, drop_last=drop_last) \\\n",
    "            .unbatched().shuffle(1024).batched(bs)\n",
    "    else:\n",
    "        train_loader = DataLoader(train, batch_size=bs, num_workers=dl_workers, pin_memory=True, drop_last=drop_last, shuffle=True)\n",
    "        val_loader = DataLoader(val, batch_size=bs, num_workers=dl_workers, pin_memory=True, drop_last=drop_last)\n",
    "    \n",
    "    val_loss = torch.nan\n",
    "    avg_train_loss = torch.nan\n",
    "    \n",
    "    if hasattr(model, 'setup'):\n",
    "        model.setup(device)\n",
    "    \n",
    "    try:\n",
    "        scheduler = None\n",
    "\n",
    "        if trainable_params is None: trainable_params = model.parameters()\n",
    "        all_params = set(trainable_params)\n",
    "        customized_params = set()\n",
    "        groups = []\n",
    "        group_map = {}\n",
    "        for name,m in model.named_modules():\n",
    "            if hasattr(m, 'no_weight_decay') or hasattr(m, 'lr_scale'):\n",
    "                m_trainable = [x for x in m.parameters() if x in all_params]\n",
    "                if not m_trainable: continue\n",
    "                customized_params |= set(m_trainable)\n",
    "                m_wd = 0 if hasattr(m, 'no_weight_decay') else weight_decay\n",
    "                m_lr = lr * getattr(m, 'lr_scale', 1)\n",
    "                group = group_map.get((m_wd, m_lr), None)\n",
    "                if not group:\n",
    "                    group = {\"params\": [], \"names\": [], \"weight_decay\": m_wd, \"lr\": m_lr}\n",
    "                    groups.append(group)\n",
    "                    group_map[(m_wd, m_lr)] = group\n",
    "                group['params'] += m_trainable\n",
    "                group['names'].append(name)\n",
    "                \n",
    "        other_params = all_params - customized_params\n",
    "        \n",
    "        if other_params:\n",
    "            groups = groups + [\n",
    "                {\"names\": [\"other\"], \"params\": list(other_params), \"weight_decay\": weight_decay },\n",
    "            ]\n",
    "\n",
    "        optimizer = torch.optim.AdamW(lr=lr, betas=(0.9, 0.95), fused=device!='cpu', params=groups)\n",
    "        model._optimizer = optimizer\n",
    "        scaler = torch.cuda.amp.GradScaler(enabled=half)\n",
    "        scheduler = torch.optim.lr_scheduler.OneCycleLR(\n",
    "            optimizer, pct_start=pct_start, steps_per_epoch=math.ceil(train.total_samples/bs), epochs=epochs,\n",
    "            max_lr=[pg.get('lr', lr) for pg in groups],\n",
    "            final_div_factor=25)\n",
    "        \n",
    "        it = 0\n",
    "        next_val_it = it + 50\n",
    "        next_chkpt_it = chkpt_every_iters\n",
    "        next_table_it = table_row_every_iters\n",
    "                \n",
    "        visual.show()\n",
    "\n",
    "        running_loss = [0]\n",
    "        \n",
    "        for epoch in mb:\n",
    "            bar = progress_bar(train_loader, total=train.total_samples//bs, parent=mb)\n",
    "            for args in bar:\n",
    "                with record_function(\"forward\"):\n",
    "                    args = [x.to(device, non_blocking=True) for x in args]\n",
    "\n",
    "                    # zero the parameter gradients\n",
    "                    optimizer.zero_grad(set_to_none=True)\n",
    "\n",
    "                    with torch.autocast(device_type=device, dtype=torch.float16 if half else torch.float32, enabled=device!='cpu'):\n",
    "                        ps, loss = model(*args)\n",
    "                    loss = loss.mean()\n",
    "\n",
    "                with record_function(\"backward\"):\n",
    "                    scaler.scale(loss).backward()\n",
    "\n",
    "                    if clip_gradient_norm:\n",
    "                        scaler.unscale_(optimizer)\n",
    "                        # Since the gradients of optimizer's assigned params are unscaled, clips as usual:\n",
    "                        torch.nn.utils.clip_grad_norm_(model.parameters(), clip_gradient_norm)\n",
    "\n",
    "                    scaler.step(optimizer)\n",
    "                    scaler.update()\n",
    "\n",
    "                    scheduler.step()\n",
    "\n",
    "                    if profiler is not None: profiler.step()\n",
    "\n",
    "                with record_function(\"running_loss\"):\n",
    "                    running_loss.append(loss.item())\n",
    "                    running_loss = running_loss[-5:]\n",
    "                    avg_train_loss = sum(running_loss)/len(running_loss)\n",
    "\n",
    "                if it >= next_chkpt_it:\n",
    "                    with record_function(\"checkpoint\"):\n",
    "                        next_chkpt_it += chkpt_every_iters\n",
    "                        torch.save(model.state_dict(), f'{checkpoint_path}/{it:08d}.pt')\n",
    "                    \n",
    "                if it >= next_val_it:\n",
    "                    next_val_it += run_valid_every_iters\n",
    "                    with record_function(\"validation\"):\n",
    "                        with record_function(\"model.eval\"):\n",
    "                            model.eval()\n",
    "                        with torch.no_grad():\n",
    "                            val_loss = 0\n",
    "                            val_samples = 0\n",
    "                            for args in val_loader:\n",
    "                                args = [x.to(device, non_blocking=True) for x in args]\n",
    "                                with torch.autocast(device_type=device, dtype=torch.float16 if half else torch.float32, enabled=device!='cpu'):\n",
    "                                    ps, loss = model(*args)\n",
    "                                N = args[0].shape[0]\n",
    "                                val_loss += loss.mean().item() * N\n",
    "                                val_samples += N\n",
    "                            val_loss = val_loss / val_samples\n",
    "                        with record_function(\"model.train\"):\n",
    "                            model.train()\n",
    "                    with record_function(\"plotting\"):\n",
    "                        visual.add_data(it, scheduler.get_last_lr(), avg_train_loss, val_loss)\n",
    "                \n",
    "                if it >= next_table_it:\n",
    "                    visual.add_table_row(it, avg_train_loss, val_loss)\n",
    "                    next_table_it += table_row_every_iters\n",
    "\n",
    "                it += bs\n",
    "                visual.on_iter(bar, it, avg_train_loss, val_loss)\n",
    "    except KeyboardInterrupt:\n",
    "        mb.write(f\"interrupted\")\n",
    "        mb.show()\n",
    "        pass\n",
    "    finally:\n",
    "        visual.add_table_row(it, avg_train_loss, val_loss)\n",
    "        mb.show()\n",
    "        visual.hide()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00406652",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "import nbdev; nbdev.nbdev_export()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ffffe92",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
